package com.milvus.imagesearch;

import com.milvus.featureExtractor.ImageFeatureExtractor;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.io.File;
import java.io.IOException;

public class Main {
    private static final String COLLECTION_NAME = "image_collection";
    public static void main(String[] args) throws IOException {
        ImageFeatureExtractor extractor = new ImageFeatureExtractor();
        MilvusManager milvusManager = new MilvusManager();
        ImageSearcher searcher = new ImageSearcher();

        // 创建Milvus集合
        milvusManager.createCollection(COLLECTION_NAME);

        // 假设有一个图像文件列表
        File[] imageFiles = new File("C:\\Users\\瑾瑜\\Desktop\\sql\\20230621").listFiles();
        if (imageFiles != null) {
            for (int i = 0; i < imageFiles.length; i++) {
                INDArray features = extractor.extractFeatures(imageFiles[i]);
                milvusManager.insertData( COLLECTION_NAME,features,imageFiles[i].getName());
            }
        }

        System.out.println("插入完成");

        // 查询图像
        File queryImage = new File("C:\\Users\\瑾瑜\\Desktop\\sql\\20230621\\AD1C2548.jpg");
        INDArray queryFeatures = extractor.extractFeatures(queryImage);
        searcher.search(COLLECTION_NAME,queryFeatures);
    }
}
